package org.dwpd.controller.ai;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.videosynthesis.VideoSynthesis;
import com.alibaba.dashscope.aigc.videosynthesis.VideoSynthesisParam;
import com.alibaba.dashscope.aigc.videosynthesis.VideoSynthesisResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.community.model.dashscope.WanxImageModel;
import dev.langchain4j.data.image.Image;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.output.Response;
import org.dwpd.assistant.MemoryChatAssistant;
import org.dwpd.assistant.SeparateChatAssistant;
import org.dwpd.domain.ai.ChatMessages;
import org.dwpd.utils.MongoChatMemoryStore;
import org.dwpd.utils.UserHolder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

@RestController
@RequestMapping("/ai")
public class ChatController {

    @Autowired
    QwenChatModel qwenChatModel;

    @Autowired
    MemoryChatAssistant memoryChatAssistant;

    @Autowired
    SeparateChatAssistant separateChatAssistant;

    @Autowired
    MongoChatMemoryStore mongoChatMemoryStore;

    WanxImageModel wanxImageModel = WanxImageModel.builder()
            .modelName("wanx2.1-t2i-plus")
            .apiKey(System.getenv("DASH_SCOPE_API_KEY"))
            .build();

    /**
     * 记忆聊天-优先使用分离记忆聊天
     * @param message 提示词
     */
    @PostMapping("/chat")
    public String model(@RequestParam(value = "message", defaultValue = "Hello") String message) throws NoApiKeyException, InputRequiredException {

        String chatJudge = qwenChatModel.chat("对下述内容进行判断，如果是要求生成一张图片，则返回数字1;如果是要求生成一个视频，则返回数字2;以上情况都不符合则返回数字0\n" + message);
        if (chatJudge.equals("1")) {
            Response<Image> response = wanxImageModel.generate(message);

            return response.content().url().toString();
        }
        else if (chatJudge.equals("2")) {
            return getVideo("0", message);
        }
//        return qwenChatModel.chat(message);
        return memoryChatAssistant.chat(message);
    }

    /**
     * 分离记忆聊天
     * @param memoryId 用来区分聊天记忆
     * @param message 提示词
     */
    @PostMapping("/separateChat")
    public String separateChat(@RequestParam(value = "memoryId") String memoryId, @RequestParam(value = "message", defaultValue = "Hello") String message) throws NoApiKeyException, InputRequiredException {

        if (memoryId.equals("0")) {
            // 创建新的记忆文档
            return mongoChatMemoryStore.createNewMemory();
        }
        String chatJudge = qwenChatModel.chat("对下述内容进行判断，如果是要求生成一张图片，则返回数字1;如果是要求生成一个视频，则返回数字2;以上情况都不符合则返回数字0\n" + message);
        if (chatJudge.equals("1")) {
            return getImage(memoryId, message);
        }
        else if (chatJudge.equals("2")) {
            return getVideo(memoryId, message);
        }
        return separateChatAssistant.chat(memoryId, message);
    }

    /**
     * 联网模式
     * @param message 提示词
     */
    @PostMapping("/chatSearch")
    public String chatSearch(@RequestParam(value = "message", defaultValue = "Hello") String message) throws NoApiKeyException, InputRequiredException {
        String chatJudge = qwenChatModel.chat("对下述内容进行判断，如果是要求生成一张图片，则返回数字1，否则则返回数字0\n" + message);
        if (chatJudge.equals("1")) {
            Response<Image> response = wanxImageModel.generate(message);
            return response.content().url().toString();
        }
        return searchModel(message);
    }

    /**
     * 删除指定聊天记忆
     */
    @PostMapping("/removeMemory")
    public void removeChatMemory(String memoryId) {
        mongoChatMemoryStore.deleteMessages(memoryId);
    }

    /**
     * 根据用户ID获取聊天记忆组
     * @return
     */
    @GetMapping("/getMemoryIds")
    public List<ChatMessages> getMemoryIdsByUserId() {
        Integer userId = UserHolder.getUser().getUserId();
        if (userId == null) {
            return new ArrayList<>();
        }
        List<ChatMessages> messageList = mongoChatMemoryStore.getMessagesByUserId(userId);
        if (messageList.isEmpty()) {
            return new ArrayList<>();
        }
        return messageList;
    }

    private String searchModel(String message) throws NoApiKeyException, InputRequiredException {
        Generation gen = new Generation();
        Message systemMsg = Message.builder()
                .role(Role.SYSTEM.getValue())
                .content("你是乐拼一夏APP的专属AI，你叫小乐，能够为用户提供准确的位置距离分析，" +
                        "并且能够搜索网络上实际的数据来为用户提供打车所需要的花费区间，" +
                        "比如从广州应用科技学院肇庆校区到肇庆东站大约12公里，平均要花费15-25元，" +
                        "在列举金额时可以提供打车平台的数据如滴滴、花小猪、哈啰，同时你能够跟用户提供趣味互动")
                .build();
        Message userMsg = Message.builder()
                .role(Role.USER.getValue())
                .content(message)
                .build();
        GenerationParam param = GenerationParam.builder()
                // 若没有配置环境变量，请用百炼API Key将下行替换为：.apiKey("sk-xxx")
                .apiKey(System.getenv("DASH_SCOPE_API_KEY"))
                // 此处以qwen-plus为例，可按需更换模型名称。模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
                .model("qwen-plus")
                .messages(Arrays.asList(systemMsg, userMsg))
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .enableSearch(true)
                .build();
        return gen.call(param).getOutput().getChoices().get(0).getMessage().getContent();
    }

    // 生成图片链接
    private String getImage(String memoryId, String message) {
        Response<Image> response = wanxImageModel.generate(message);
        String imageUrl = response.content().url().toString();
        mongoChatMemoryStore.appendContentToMemory(memoryId, message, imageUrl);
        return imageUrl;
    }

    // 生成视频链接
    private String getVideo(String memoryId, String message) throws NoApiKeyException, InputRequiredException {
        VideoSynthesis vs = new VideoSynthesis();
        VideoSynthesisParam param =
                VideoSynthesisParam.builder()
                        .model("wanx2.1-t2v-turbo")
                        .prompt(message)
                        .size("1280*720")
                        .apiKey(System.getenv("DASH_SCOPE_API_KEY"))
                        .build();
        System.out.println("please wait...");
        VideoSynthesisResult videoSynthesisResult = vs.call(param);
        String videoUrl = videoSynthesisResult.getOutput().getVideoUrl();
        mongoChatMemoryStore.appendContentToMemory(memoryId, message, videoUrl);
        return videoUrl;
    }
}
